/*
 * 版权所有 (C) 2015 知启蒙(ZHIQIM) 保留所有权利。[遇见知启蒙，邂逅框架梦]
 * 
 * https://zhiqim.org/project/zhiqim_framework/zhiqim_kernel.htm
 *
 * Zhiqim Kernel is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */
package org.zhiqim.kernel.model.codes;

import java.io.ByteArrayOutputStream;
import java.nio.charset.Charset;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.Cipher;

import org.zhiqim.kernel.constants.CodeConstants;
import org.zhiqim.kernel.model.objects.KV;
import org.zhiqim.kernel.model.objects.SS;
import org.zhiqim.kernel.util.Asserts;
import org.zhiqim.kernel.util.Hexs;

/**
 * RSA非对称加解密
 *
 * @version v1.0.0 @author zouzhigang 2018-8-22 新建与整理
 */
public class RSA implements CodeConstants
{
    /**********************************************************************************/
    //生成密钥对
    /**********************************************************************************/
    
    /**
     * 生成一对公钥和私钥，第一个为公钥，第二个为私钥，默认1024
     */
    public static SS buildKeyPair()
    {
        return buildKeyPairBase64(KiB);
    }
    
    /**
     * 生成一对公钥和私钥
     * 
     * @param keySize   密钥长度，如512/1024/2048
     * @return          第一个为公钥，第二个为私钥
     */
    public static SS buildKeyPairBase64(int keySize)
    {
        try
        {
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(_RSA_);
            keyPairGen.initialize(keySize);
            KeyPair keyPair = keyPairGen.generateKeyPair();
            
            String publicKey = Base64.encode(keyPair.getPublic().getEncoded());
            String privateKey = Base64.encode(keyPair.getPrivate().getEncoded());
            return new SS(publicKey, privateKey);
        }
        catch(NoSuchAlgorithmException e)
        {
            throw Asserts.exception(e);
        }
    }
    
    /**
     * 生成一对公钥和私钥
     * 
     * @param keySize   密钥长度，如512/1024/2048
     * @return          第一个为公钥，第二个为私钥
     */
    public static KV<byte[], byte[]> buildKeyPair(int keySize)
    {
        try
        {
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(_RSA_);
            keyPairGen.initialize(keySize);
            KeyPair keyPair = keyPairGen.generateKeyPair();
            
            return new KV<>(keyPair.getPublic().getEncoded(), keyPair.getPrivate().getEncoded());
        }
        catch(NoSuchAlgorithmException e)
        {
            throw Asserts.exception(e);
        }
    }
    
    /**********************************************************************************/
    //标准加解密
    /**********************************************************************************/
    
    /**
     * 加密（使用公钥加密），支持常用的512/1024/2048/3072/4096长度适配
     * 
     * @param src           原数据
     * @param publicKey     公钥字符串（BASE64）
     * @return              加密后数据
     * @throws Exception    
     */
    public static byte[] encrypt(byte[] src, String publicKey) throws Exception
    {
        return encrypt(src, Base64.decode(publicKey));
    }
    
    /**
     * 加密（使用公钥加密），支持常用的512/1024/2048/3072/4096长度适配
     * 
     * @param src           原数据
     * @param publicKey     公钥字节数组
     * @return              加密后数据
     * @throws Exception    
     */
    public static byte[] encrypt(byte[] src, byte[] publicKey) throws Exception
    {
        int keySize = 0;
        switch (publicKey.length)
        {
        case 94: keySize = 512; break;//343
        case 162:keySize = 1024;break;//634
        case 294:keySize = 2048;break;//1217
        case 422:keySize = 3072;break;//1793
        case 550:keySize = 4096;break;//2376
        }
        
        Asserts.as(keySize != 0?null:"无法匹配密钥长度，请调用encrypt(byte[] src, byte[] publicKey, int keySize)方法");
        return encrypt(src, publicKey, keySize);
    }
    
    /**
     * 加密（使用公钥加密）
     * 
     * @param src           原数据
     * @param publicKey     公钥字节数组
     * @param keySize       密钥长度，如512/1024/2048
     * @return              加密后数据
     * @throws Exception    可能的异常
     */
    public static byte[] encrypt(byte[] src, byte[] publicKey, int keySize) throws Exception
    {
        //1.组装公钥
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKey);
        KeyFactory keyFactory = KeyFactory.getInstance(_RSA_);
        PublicKey key = keyFactory.generatePublic(keySpec);
        
        //2.初始化密码器
        Cipher cipher = Cipher.getInstance(_RSA_);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        
        //3.加密时分块加密，每块支持大小为keySize / 8 - 11
        int length = src.length;
        int chunkSize = keySize / 8 - 11;
        
        for (int i=0;i*chunkSize < length;i++)
        {
            int count = length > (i+1) * chunkSize?chunkSize:length - i*chunkSize;
            out.write(cipher.doFinal(src, i * chunkSize, count));
        }
        
        return out.toByteArray();
    }
    
    /**
     * 解密（使用私钥加密），支持常用的512/1024/2048/3072/4096长度适配
     * 
     * @param src           原数据
     * @param privateKey    私钥字符串（BASE64）
     * @return              加密后数据
     * @throws Exception    可能的异常
     */
    public static byte[] decrypt(byte[] src, String privateKey) throws Exception
    {
        return decrypt(src, Base64.decode(privateKey));
    }
    
    /**
     * 解密（使用私钥加密），支持常用的512/1024/2048/3072/4096长度适配
     * 
     * @param src           原数据
     * @param privateKey    私钥字节数组
     * @return              加密后数据
     * @throws Exception    可能的异常
     */
    public static byte[] decrypt(byte[] src, byte[] privateKey) throws Exception
    {
        int keyLen = privateKey.length;
        int keySize = 0;
        
        if (keyLen >= 330 && keyLen <= 360)//342-347
            keySize = 512;
        else if (keyLen >= 620 && keyLen <= 650)//631-637
            keySize = 1024;
        else if (keyLen >= 1200 && keyLen <= 1230)//1214-1220
            keySize = 2048;
        else if (keyLen >= 1780 && keyLen <= 1710)//1792-1796
            keySize = 3072;
        else if (keyLen >= 2360 && keyLen <= 2390)//2373-2376
            keySize = 4096;
        
        Asserts.as(keySize != 0?null:"无法匹配密钥长度，请调用decrypt(byte[] src, byte[] keyBytes, int keySize)方法");
        return decrypt(src, privateKey, keySize);
    }

    /**
     * 解密（使用私钥加密）
     * 
     * @param src           原数据
     * @param privateKey    私钥字节数组
     * @param keySize       密钥长度，如512/1024/2048
     * @return              加密后数据
     * @throws Exception    可能的异常
     */
    public static byte[] decrypt(byte[] src, byte[] privateKey, int keySize) throws Exception
    {
        //1.组装私钥
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKey);
        KeyFactory keyFactory = KeyFactory.getInstance(_RSA_);
        PrivateKey key = keyFactory.generatePrivate(keySpec);
        
        //2.初始化密码器
        Cipher cipher = Cipher.getInstance(_RSA_);
        cipher.init(Cipher.DECRYPT_MODE, key);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        
        //3.解密时分块加密，每块支持大小为keySize / 8
        int length = src.length;
        int chunkSize = keySize / 8;
        
        for (int i=0;i*chunkSize < length;i++)
        {
            int count = length > (i+1) * chunkSize?chunkSize:length - i*chunkSize;
            out.write(cipher.doFinal(src, i*chunkSize, count));
        }
        
        return out.toByteArray();
    }
    
    /**********************************************************************************/
    //字符串加解密
    /**********************************************************************************/
    
    /** 字符串加密 */
    public static byte[] encrypt(String src, Charset charset, String publicKey) throws Exception
    {
        return encrypt(src.getBytes(charset), publicKey);
    }
    
    /** 字符串加密 */
    public static byte[] encrypt(String src, String encoding, String publicKey) throws Exception
    {
        return encrypt(src, Charset.forName(encoding), publicKey);
    }
    
    /** 字符串解密 */
    public static String decrypt(byte[] src, String privateKey, Charset charset) throws Exception
    {
        return new String(decrypt(src, privateKey), charset);
    }
    
    /** 字符串解密 */
    public static String decrypt(byte[] src, String privateKey, String encoding) throws Exception
    {
        return decrypt(src, privateKey, Charset.forName(encoding));
    }
    
    /**********************************************************************************/
    //字符串BASE64加解密
    /**********************************************************************************/
    
    /** 字符串加密 */
    public static String encryptBase64UTF8(String src, String publicKey) throws Exception
    {
        return encryptBase64(src, _UTF_8_C_, publicKey);
    }
    
    /** 字符串加密 */
    public static String encryptBase64(String src, String encoding, String publicKey) throws Exception
    {
        return encryptBase64(src, Charset.forName(encoding), publicKey);
    }
    
    /** 字符串加密 */
    public static String encryptBase64(String src, Charset charset, String publicKey) throws Exception
    {
        return Base64.encode(encrypt(src.getBytes(charset), publicKey));
    }

    /** 字符串解密 */
    public static String decryptBase64UTF8(String src, String privateKey) throws Exception
    {
        return decryptBase64(src, privateKey, _UTF_8_C_);
    }
    
    /** 字符串解密 */
    public static String decryptBase64(String src, String privateKey, String encoding) throws Exception
    {
        return decryptBase64(src, privateKey, Charset.forName(encoding));
    }
    
    /** 字符串解密 */
    public static String decryptBase64(String src, String privateKey, Charset charset) throws Exception
    {
        byte[] bytes = Base64.decode(src);
        return new String(decrypt(bytes, privateKey), charset);
    }
    /**********************************************************************************/
    //字符串HEX加解密
    /**********************************************************************************/
    
    /** 字符串加密 */
    public static String encryptHexUTF8(String src, String publicKey) throws Exception
    {
        return encryptHex(src, _UTF_8_C_, publicKey);
    }
    
    /** 字符串加密 */
    public static String encryptHex64(String src, String encoding, String publicKey) throws Exception
    {
        return encryptHex(src, Charset.forName(encoding), publicKey);
    }
    
    /** 字符串加密 */
    public static String encryptHex(String src, Charset charset, String publicKey) throws Exception
    {
        return Hexs.toHexString(encrypt(src.getBytes(charset), publicKey));
    }

    /** 字符串解密 */
    public static String decryptHexUTF8(String src, String privateKey) throws Exception
    {
        return decryptHex(src, privateKey, _UTF_8_C_);
    }
    
    /** 字符串解密 */
    public static String decryptHex(String src, String privateKey, String encoding) throws Exception
    {
        return decryptHex(src, privateKey, Charset.forName(encoding));
    }
    
    /** 字符串解密 */
    public static String decryptHex(String src, String privateKey, Charset charset) throws Exception
    {
        byte[] bytes = Hexs.toBytes(src);
        return new String(decrypt(bytes, privateKey), charset);
    }
}
